rm(list = ls())

library(data.table)

load('./data/panel_month_dummies.RData')
panel <- panel[!is.na(share_asian_immigrant)]

permute <- function(permutation_data, n_reps) {
  results <- vector(mode='list', length=n_reps)
  for (r in seq(n_reps)) {
    #print(paste("Rep", r))
    true_treatments <- data.table(permutation_data)[, .(is_treated=unique(losing_income_d)==1), by=panelvar]
    
    permutation_treatments <- data.table(
      is_treated_permutation = sample(true_treatments$is_treated),
      permutation_id = unique(permutation_data$panelvar)
    )
    
    permutation_data$is_treated_permutation <- NULL
    permutation_data <- merge(permutation_data, permutation_treatments, by.x='panelvar', by.y='permutation_id')
    
    nbzi_b_c <- pscl::zeroinfl(asian_hc ~ 1 | share_asian_immigrant,
                               data=permutation_data[covid2==0 & is_treated_permutation==0], dist = 'negbin', link = 'logit')
    nbzi_b_t <- pscl::zeroinfl(asian_hc ~ 1 | share_asian_immigrant,
                               data=permutation_data[covid2==0 & is_treated_permutation==1], dist = 'negbin', link = 'logit')
    nbzi_a_c <- pscl::zeroinfl(asian_hc ~ 1 | share_asian_immigrant,
                               data=permutation_data[covid2==1 & is_treated_permutation==0], dist = 'negbin', link = 'logit')
    nbzi_a_t <- pscl::zeroinfl(asian_hc ~ 1 | share_asian_immigrant,
                               data=permutation_data[covid2==1 & is_treated_permutation==1], dist = 'negbin', link = 'logit')
    
    
    atet <- mean((predict(nbzi_a_t, permutation_data[covid2==1 & is_treated_permutation==1], type = 'response') -
                    predict(nbzi_b_t, permutation_data[covid2==1 & is_treated_permutation==1], type = 'response')) -
                   (predict(nbzi_a_c, permutation_data[covid2==1 & is_treated_permutation==1], type = 'response') -
                      predict(nbzi_b_c, permutation_data[covid2==1 & is_treated_permutation==1], type = 'response')))
    results[[r]] <- atet
  }
  return(results)
}


set.seed(879)
rr <- permute(panel, 500)

save(rr, file='./output/zi_neg_binomial_ri_tmp.RData')